import time
import datetime
from torch.cuda.amp import GradScaler

from clearml import Task, Logger

import os
import torch

import sys

# 加载数据就和处理


# 加载其他工具类：保存函数、写日志函数、损失函数等等
from utils.loops import train, evaluate, test
from utils.opti import get_opti, get_loss_fn
from utils.opti import scheduler_update
from utils.utils import load_checkpoint, random_seed, save_checkpoint, get_args
from model.getmodel import get_model

def main():
    # 获取配置信息
    config = get_args()
    print(config)
    if config.cm:
        task = Task.init(project_name=config.taskproject, task_name=config.name)  # 负责记录输入输出
        task.connect(config)
        logger = Logger.current_logger()  # 显式报告 包括标量图、线图、直方图、混淆矩阵、2D 和 3D 散点图、文本日志记录、表格以及图像上传和报告
    # 定义本次输出的根目录和checkpoint目录
    checkpoint_path = os.path.join(config.outputs,
                                   config.name,
                                   datetime.datetime.now().strftime("%Y-%m-%d-%H.%M.%S"),
                                   'checkpoints')
    if not os.path.exists(checkpoint_path):  # 创建最深的路径，上级路径也就有了
        os.makedirs(checkpoint_path)

    # 随机种子
    if config.seed is not None:  # config.seed填写(包括0)即可执行，不填写则为None
        random_seed(config.seed)

    device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
    print("device:", device)


    if config.data_type == "fer2013":
        from data.csvdataloaders import get_csvdataloaders
        train_loader, val_loader, test_loader = get_csvdataloaders(path=config.data_path,
                                                                   bs=config.batch_size,
                                                                   num_workers=config.num_workers,
                                                                   augment=True)
    else:       # 默认ImageNet数据加载器
        from data.imagenet import get_imagenet_dataloader
        train_loader, val_loader, test_loader = get_imagenet_dataloader(path=config.data_path,
                                                                        bs=config.batch_size,
                                                                        num_workers=config.num_workers,
                                                                        augment=True)

    model, net_path = get_model(arch=config.arch, device=device)
    loss_fn = get_loss_fn(config)    # 获取损失函数

    if config.cm:
        task.upload_artifact(name="net path", artifact_object=net_path)  # 记录定义net的py文件

    if config.mode=="test":
        model = load_checkpoint(config.best_checkpoint, model=model)
        print("加载checkpoint成功：", config.best_checkpoint)
        test(model, test_loader, Ncrop=config.Ncrop, device=device, detail=True)
        sys.exit()


    scaler = GradScaler() if config.amp else False  # 是否使用amp
    optimizer, scheduler = get_opti(config, model)

    print("Epoch \t Time \t Train Loss \t Train ACC \t Val Loss \t Val ACC \t Test Acc")
    best_acc = 0
    for epoch in range(1, config.epochs + 1):
        start_t = time.time()

        # 训练
        train_loss, train_acc = train(model, train_loader, loss_fn, optimizer, scaler=scaler, config=config, device=device)
        val_loss, val_acc = evaluate(model, val_loader, config=config, device=device)
        test_acc = test(model, test_loader, Ncrop=config.Ncrop, device=device, detail=False)

        # 更新lr衰减
        scheduler=scheduler_update(config, scheduler, val_acc=val_acc)

        # 上传信息
        if config.cm:
            logger.report_scalar(title='Loss', series='Train', value=train_loss, iteration=epoch)
            logger.report_scalar(title='Loss', series='Val', value=val_loss, iteration=epoch)
            logger.report_scalar(title='Accuracy', series='Train', value=train_acc, iteration=epoch)
            logger.report_scalar(title='Accuracy', series='Val', value=val_acc, iteration=epoch)
            logger.report_scalar(title='Test Accuracy', series='Test', value=test_acc, iteration=epoch)
            lr = optimizer.state_dict()['param_groups'][0]['lr']
            logger.report_scalar(title='lr', series='epoch', value=lr, iteration=epoch)

        # 如果出现最佳Acc
        note = ""
        if val_acc > best_acc:
            "如果出现最佳model：上传BestTestAcc-->保存模型（按需）-->设置note提醒字符"
            best_acc = val_acc
            if config.cm:
                logger.report_scalar(title='Best Accuracy', series="Val", value=best_acc, iteration=epoch)
                task.connect({"Test Acc": test_acc})
            if config.save_checkpoint:
                save_checkpoint(state={"model_state_dict": model.state_dict(), "epoch": epoch},
                                best_filename=os.path.join(checkpoint_path, f"best_checkpoint_{str(best_acc)}.pth"))
            note = "saves best"

        # 每Epoch的控制台信息
        print("%d\t %.4f \t %.4f \t %.4f \t %.4f \t %.4f \t %.4f \t %s" % (
            epoch,
            time.time() - start_t,
            train_loss,
            train_acc,
            val_loss,
            val_acc,
            test_acc,
            note))




if __name__ == '__main__':
    main()
